import json
import logging
import threading
import time

from tools.algo.config import AllConfig, AlgoConfig, AlgoRunConfig
from tools.algo.task_schedule import TaskScheduler

from celery_worker import app
from tools.constant.config_key import CAMERA_KEY, CAMERA_STREAM, LIVE_HOST, LIVE_PORT
from tools.constant.redis_keys import *
from tools.utils.logging_setting import set_logging
from tools.utils.redis_client import redis_client_instance

set_logging()

client = redis_client_instance()
pubsub = client.pubsub()

all_config= AllConfig()
algo_config  = AlgoConfig()
taskSchedule = TaskScheduler()
algo_run_config = AlgoRunConfig()

def thread_process(algo_name,same_algos=None):
    # 调佣任务，分配桶子
    while True:
        if algo_run_config.get_stop_algo(algo_name):
            taskSchedule.task_clean(algo_name)
            exit()
        algo_size = taskSchedule.task_size(algo_name)
        logging.debug(f"算法--{algo_name},当前队列为{algo_size}")
        start_algo = False
        if algo_size>=0 and algo_size <= 100:
            start_algo = True
        else:
            start_algo = False
        if start_algo:
            # TODO 兼顾多个的情况
            load_model,run_model= algo_run_config.need_to_load_model(algo_name)
            # load_model=True
            if load_model == run_model and not load_model and algo_name not in algo_run_config.white_list_algo:
                logging.debug(f"算法休眠--{algo_name}")
                time.sleep(30)
                continue
            algo_configs_map = algo_config.get_algo_configs(algo_name,same_algos)
            qus = algo_run_config.get_algo_router(algo_name)
            if algo_configs_map and len(algo_configs_map):
                # 有摄像头配置
                for cam,cf in algo_configs_map.items():
                    if cf and len(cf) and isinstance(cf,str):
                        cam_conf = json.loads(cf)
                    else:
                        cam_conf = {}
                    cam_conf[CAMERA_KEY]=f"camera:frames:{cam}"
                    if load_model:
                        task = run_algo_with_config(algo_name, cam_conf, qus)
                        end_ts = int(time.time())
                        same_algos = algo_run_config.get_same_algo(algo_name)
                        if same_algos and len(same_algos)>0:
                            for same_algo in same_algos:
                                taskSchedule.add_task(same_algo, task.task_id, ts=end_ts)
                        else:
                            taskSchedule.add_task(algo_name, task.task_id,ts=end_ts)
            elif algo_name in  algo_run_config.white_list_algo:
                # 白名单，直接调用
                if load_model:
                    task = run_algo_with_config(algo_name, {}, qus,send_without_config=True)
                    end_ts = int(time.time())
                    taskSchedule.add_task(algo_name, task.task_id,ts=end_ts)
            #判断是否是首次加载模型
            if load_model != run_model:
                #加载模型，延期10秒
                time.sleep(10)
            else:
                time.sleep(all_config.interval_seconds)
            # #摄像头配置
        time.sleep(1)

def thread_process_type3(algo_name):
    # 第三类算法，直接调用一次，持续产生数据，每60秒检查一次，直到收到退出信号
    last_model_status = False
    while True:
        if algo_run_config.get_stop_algo(algo_name):
            exit()
        start_algo = True
        if start_algo:
            load_model,run_model= algo_run_config.need_to_load_model(algo_name)
            if load_model == run_model and not load_model and algo_name not in algo_run_config.white_list_algo:
                logging.debug(f"算法休眠--{algo_name}")
                time.sleep(30)
                continue
            algo_configs_map = algo_config.get_algo_configs(algo_name)
            qus = algo_run_config.get_algo_router(algo_name)
            last_model_status =  (load_model !=last_model_status)
            if algo_configs_map and len(algo_configs_map):
                # 有摄像头配置
                for cam,cf in algo_configs_map.items():
                    if cf and len(cf) and isinstance(cf,str):
                        cam_conf = json.loads(cf)
                    else:
                        cam_conf = {}
                    cam_conf[CAMERA_KEY]=f"camera:frames:{cam}"
                    camera_info = client.hget(REDIS_CAMERA_INFO_KEY, cam)
                    if camera_info:
                        cam_conf[CAMERA_STREAM] = camera_info.decode("utf-8")
                    cam_conf["RUN"] = load_model
                    cam_conf["RTMP_URL"] ="rtmp://172.16.30.62:1935/live/8492"
                    # cam_conf["RTMP_URL"] = f"rtmp://{LIVE_HOST}:{LIVE_PORT}/{cam}"
                    if last_model_status:
                        run_algo_with_config(algo_name, cam_conf, qus)
            else:
                run_algo_with_config(algo_name, {}, qus)
                exit()
        time.sleep(60)




def run_algo_with_config(algo_name,config,qus,send_without_config=False):
    if isinstance(qus,set):
        qus = ",".join(qus)
    if send_without_config:
        result = app.send_task(algo_name,args=[config], queue=qus)
    else:
        result = app.send_task(algo_name, args=[config],queue=qus)
    return result


def handle_task_result(algo_name):
    while True:
        if algo_name not in algo_run_config.algo_type2:
            taskSchedule.handler_task(algo_name,client=client)
        else:
            same_algos = algo_run_config.get_same_algo(algo_name)
            taskSchedule.handler_task_same_model(same_algos, algo_name, algo_config,client=client)
        time.sleep(0.5)

# def handle_task_result_model(algo_name,parent,qus):
#     while True:
#         taskSchedule.handler_task_same_model(algo_name,parent,qus,client=client)
#         time.sleep(0.5)

all_algorithm_list = []
all_algorithm_thread_list = []
all_algorithm_thread_reulst_list = []


def start_check_thead_status():

    online_alg=threading.Thread(target=thread_registed_task, args=(),
                                          name="check_online_algo")
    read_redis_thread = threading.Thread(target=read_redis_config, args=(),
                                          name="read_redis_config")
    online_alg.start()
    read_redis_thread.start()
    while True:
        algorithm_list = algo_run_config.get_all_algo_list()
        current_list = []
        current_result_list = []
        for algo_name in algorithm_list:
            algo_thread_name = algo_name + "-thread"
            algo_thread_reult_name = algo_name + "-result"
            current_list.append(algo_thread_name)
            current_result_list.append(algo_thread_reult_name)
            # 如果有新的，增加进去
            if algo_thread_name in all_algorithm_list:
                continue
            else:
                all_algorithm_list.append(algo_thread_name)

                if algo_name in algo_run_config.algo_type3:
                    # 视频算法，启动一次就好
                    thread = threading.Thread(target=thread_process_type3, args=(algo_name,),
                                              name=algo_thread_name)
                elif algo_name in algo_run_config.algo_type2:
                    # 视频算法，启动一次就好
                    same_algo = algo_run_config.get_same_algo(algo_name)
                    thread = threading.Thread(target=thread_process, args=(algo_name,same_algo),
                                              name=algo_thread_name)
                else:
                    thread = threading.Thread(target=thread_process, args=(algo_name,),
                                          name=algo_thread_name)
                thread_result = threading.Thread(target=handle_task_result, args=(algo_name,),
                                                 name=algo_thread_reult_name)

                all_algorithm_thread_reulst_list.append(thread_result)
                all_algorithm_thread_list.append(thread)
                logging.debug(f"线程启动中 {algo_thread_name}")
                thread.start()
                time.sleep(3)
                thread_result.start()
        # 检查存活
        for t in all_algorithm_thread_list:
            if t.name not in current_list:
                # stop_algo.pop()
                algo_run_config.set_stop_algo(algo_name, True)
                continue
            if not t.is_alive():
                all_algorithm_list.remove(t.name)
                all_algorithm_thread_list.remove(t)
            logging.debug("线程%s状态-》%s " % (t.name, t.is_alive()))

        for ts in all_algorithm_thread_reulst_list:
            if not ts.is_alive():
                current_result_list.remove(ts.name)
                all_algorithm_thread_reulst_list.remove(ts)
            logging.debug("线程%s状态-》%s " % (ts.name, ts.is_alive()))
        time.sleep(5)


def get_registered_tasks():
    """
    Retrieve and update the list of registered tasks and their associated queues.

    Returns:
        set: A set of all registered tasks.
    """
    logging.info("------------------------重新检查算法列表-----------------------")
    try:
        algo_queues = inspect_workers(None)
        # algo_queues = inspect_workers()
        all_tassk_run = client.hgetall(ALGO_RUN_MODEL)
        all_tassk = client.hgetall(ALGO_LOAD_MODEL)
        if len(all_tassk) != len(all_tassk_run):
            for k,v in all_tassk.items():
                client.hset(ALGO_RUN_MODEL, k, 0)
            all_tassk_run = client.hgetall(ALGO_RUNNABLE_KEY)
        online_alg=[]
        for k,v in algo_queues.items():
            algo_run_config.algo_router[k] = v
            online_alg.append(k)
            if k.encode("utf-8") not in all_tassk.keys():
                client.hset(ALGO_LOAD_MODEL,k,0)
            if k.encode("utf-8") not in all_tassk_run.keys():
                client.hset(ALGO_RUNNABLE_KEY, k, 0)

        update_algo_config(online_alg)
    except Exception as e:
        logging.error(e)


def thread_registed_task():
    while True:
        try:
            get_registered_tasks()
        except Exception as e:
            logging.error(e)
        time.sleep(10)

def inspect_workers(search=None):
    """
    Inspect Celery workers to get their registered tasks and active queues.

    Returns:
        dict: A dictionary mapping each task to the set of queues it listens to.
    """
    load_json = {}
    worker_tasks = {}
    worker_queues = {}

    # Create an inspector object
    inspector = app.control.inspect()

    # Get all registered tasks and active queues
    registered_tasks = inspector.registered()
    active_queues = inspector.active_queues()

    if registered_tasks and active_queues:
        for worker in registered_tasks:
            # Collect tasks registered by the worker
            worker_tasks[worker] = set(registered_tasks[worker])

            # Collect queues the worker is listening to
            worker_queues[worker] = set(queue['name'] for queue in active_queues.get(worker, []))

        # Map each task to the set of queues it listens to
        for worker, tasks in worker_tasks.items():
            for task in tasks:
                if search and  not task.__contains__(search):
                    continue
                if task in load_json:
                    load_json[task].update(worker_queues[worker])
                else:
                    load_json[task] = worker_queues[worker]



    return load_json


def update_algo_config(exist_algos=[]):
    ites = client.hgetall(ALGO_CONFIG_KEY)
    algo_run_config.algo_type1=[]
    algo_run_config.algo_type2= []
    algo_run_config.algo_type3 = []
    for k, v in ites.items():
        if k.decode("utf-8") not in exist_algos:
            continue
        conf = json.loads(v.decode("utf-8"))
        if conf.__contains__("algo_type"):
            if conf["algo_type"] == 0:
                if not algo_run_config.algo_type1.__contains__(k.decode("utf-8")):
                    algo_run_config.algo_type1.append(k.decode("utf-8"))
            if conf["algo_type"] == 1:
                if not algo_run_config.algo_type2.__contains__(k.decode("utf-8")):
                    algo_run_config.algo_type2.append(k.decode("utf-8"))
                if conf.__contains__("model_id"):
                    if algo_run_config.algo_combine_algo.__contains__(conf["model_id"]):
                        if not algo_run_config.algo_combine_algo[conf["model_id"]].__contains__(k.decode("utf-8")):
                            algo_run_config.algo_combine_algo[conf["model_id"]].append(k.decode("utf-8"))
                            if k.decode("utf-8") in all_algorithm_list:
                                algo_run_config.set_stop_algo(k.decode("utf-8"), True)
                                all_algorithm_list.remove(k.decode("utf-8"))
                    else:
                        algo_run_config.algo_combine_algo[conf["model_id"]] = [k.decode("utf-8")]
            if conf["algo_type"] == 2:
                if not algo_run_config.algo_type3.__contains__(k.decode("utf-8")):
                    algo_run_config.algo_type3.append(k.decode("utf-8"))
                    if k.decode("utf-8") in all_algorithm_list:
                        algo_run_config.set_stop_algo(k.decode("utf-8"), True)
                        all_algorithm_list.remove(k.decode("utf-8"))
        else:
            if not algo_run_config.algo_type1.__contains__(k.decode("utf-8")):
                algo_run_config.algo_type1.append(k.decode("utf-8"))
    for ks in exist_algos:
        if client.hget(ALGO_CONFIG_KEY,ks) is None:
            client.hset(ALGO_CONFIG_KEY, ks, json.dumps({"@class": "java.lang.Object"}))

def read_redis_config():
    while True:
        all_algo = []
        configs = client.hgetall(ALGO_RUNNABLE_KEY)
        for k, v in configs.items():
            all_algo.append(k.decode("utf-8"))
            camera_configs = client.hgetall(AIGO_CAMERA_CONFIG_PRE+k.decode("utf-8"))
            camera_list= {}
            for cam,cf in camera_configs.items():
                try:
                    camera_list[cam.decode("utf-8")]=json.loads(cf.decode("utf-8"))
                except Exception as e:
                    camera_list[cam.decode("utf-8")] = {}
            algo_config.set_algo_configs(k.decode("utf-8"),camera_list)
        load_model_map = {}
        load_models = client.hgetall(ALGO_LOAD_MODEL)
        for k, v in load_models.items():
            load_model_map[k.decode("utf-8")] = int(v.decode("utf-8"))
        run_model_map = {}
        run_models = client.hgetall(ALGO_RUNNABLE_KEY)
        for k, v in run_models.items():
            run_model_map[k.decode("utf-8")] = int(v.decode("utf-8"))

        for k in all_algo:
            if load_model_map.__contains__(k) and run_model_map.__contains__(k):
                # set_need_to_load_model(k,load_model_map[k]>0)
                algo_run_config.set_need_to_load_model(k, (load_model_map[k] > 0, run_model_map[k] > 0))
            else:
                algo_run_config.set_need_to_load_model(k, (False,False))
        all_config.set_interval_seconds(client=client)
        time.sleep(30)


if __name__ == "__main__":
    # pass
    # print(inspect_workers("SEGMENTATION"))
    # print(redis_client_instance().llen("algo_channel"))
    # print(get_registered_tasks())
    # print(get_registered_tasks())
    start_check_thead_status()
    # online_alg = threading.Thread(target=thread_registed_task, args=(),
    #                               name="check_online_algo")
    # read_redis_thread = threading.Thread(target=read_redis_config, args=(),
    #                                      name="read_redis_config")
    # online_alg.start()
    # read_redis_thread.start()
    # print(inspect_workers())

    # thread_process_type3("OBJECT_COUNTING")
    # handle_task_result("OBJECT_SEGMENTATION")
    # st1 = threading.Thread(target=handle_task_result, args=("FACE_RETRIEVAL",))
    # st2 = threading.Thread(target=handle_task_result, args=("FACE_RETRIEVAL",))
    # st3 = threading.Thread(target=handle_task_result, args=("FACE_RETRIEVAL",))
    # st4 = threading.Thread(target=handle_task_result, args=("FACE_RETRIEVAL",))
    # st5 = threading.Thread(target=handle_task_result, args=("FACE_RETRIEVAL",))
    # st1.start()
    # st2.start()
    # st3.start()
    # st4.start()
    # st5.start()
    # get_registered_tasks()
    # print(algo_run_config.algo_combine_algo)


